package com.gitee.zhziyun.cnfsystem.utils.udputil;

import java.io.IOException;
import java.net.*;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

public class UDPReceiver extends Thread {
    private static final int BUFFER_SIZE = 1500; // 略大于最大预期包大小
    private boolean running;
    private final int port;
    private Queue<String> messageQueue;

    // 用于存储分组消息的临时缓存
    private final Map<String, MessageGroup> messageGroups = new ConcurrentHashMap<>();
    // 定时清理过期分组
    private Timer cleanupTimer = new Timer(true);

    public UDPReceiver(Queue<String> messageQueue, int port) {
        this.messageQueue = messageQueue;
        this.port = port;

        // 每5分钟清理一次过期分组 (30秒未完成的消息)
        cleanupTimer.scheduleAtFixedRate(new TimerTask() {
            @Override
            public void run() {
                long now = System.currentTimeMillis();
                messageGroups.entrySet().removeIf(entry ->
                        now - entry.getValue().lastUpdated > 30000);
            }
        }, 300000, 300000); // 5分钟检查一次
    }

    @Override
    public void run() {
        running = true;

        try (DatagramSocket socket = new DatagramSocket(port)) {
            System.out.println("UDP服务端启动，监听端口: " + port);

            byte[] buffer = new byte[BUFFER_SIZE];
            DatagramPacket packet = new DatagramPacket(buffer, buffer.length);

            while (running) {
                try {
                    socket.receive(packet);

                    // 解析分组头
                    String received = new String(
                            packet.getData(),
                            0,
                            packet.getLength(),
                            StandardCharsets.UTF_8
                    );

                    // 分组头格式: messageId|totalPackets|packetNum|data
                    int firstSep = received.indexOf('|');
                    if (firstSep == -1) continue; // 无效包

                    int secondSep = received.indexOf('|', firstSep + 1);
                    if (secondSep == -1) continue; // 无效包

                    int thirdSep = received.indexOf('|', secondSep + 1);
                    if (thirdSep == -1) continue; // 无效包

                    String messageId = received.substring(0, firstSep);
                    int totalPackets = Integer.parseInt(received.substring(firstSep + 1, secondSep));
                    int packetNum = Integer.parseInt(received.substring(secondSep + 1, thirdSep));
                    String data = received.substring(thirdSep + 1);

                    // 获取或创建消息组
                    MessageGroup group = messageGroups.computeIfAbsent(messageId,
                            k -> new MessageGroup(totalPackets));

                    // 添加数据包
                    synchronized (group) {
                        group.addPacket(packetNum, data);

                        // 检查是否完成
                        if (group.isComplete()) {
                            String completeMessage = group.getCompleteMessage();
                            messageGroups.remove(messageId);

                            System.out.printf("收到完整消息来自 %s: %s%n",
                                    packet.getAddress().getHostAddress(),
                                    completeMessage);

                            messageQueue.add(packet.getAddress().getHostAddress() + ":" + completeMessage);
                        }
                    }

                    // 重置packet长度
                    packet.setLength(buffer.length);
                } catch (IOException e) {
                    if (running) {
                        System.err.println("接收数据出错: " + e.getMessage());
                    }
                } catch (NumberFormatException e) {
                    System.err.println("无效的分组头格式");
                }
            }
        } catch (SocketException e) {
            System.err.println("Socket错误: " + e.getMessage());
        } finally {
            cleanupTimer.cancel();
        }
    }

    public void close() {
        running = false;
    }

    // 内部类用于管理分组消息
    private static class MessageGroup {
        private final String[] packets;
        private int receivedCount = 0;
        private long lastUpdated = System.currentTimeMillis();

        public MessageGroup(int totalPackets) {
            this.packets = new String[totalPackets];
        }

        public synchronized void addPacket(int packetNum, String data) {
            if (packetNum >= 0 && packetNum < packets.length && packets[packetNum] == null) {
                packets[packetNum] = data;
                receivedCount++;
                lastUpdated = System.currentTimeMillis();
            }
        }

        public synchronized boolean isComplete() {
            return receivedCount == packets.length;
        }

        public synchronized String getCompleteMessage() {
            StringBuilder sb = new StringBuilder();
            for (String packet : packets) {
                sb.append(packet);
            }
            return sb.toString();
        }
    }
}